// AssertThread.cpp : Defines the exported functions for the DLL.
//

#include "pch.h"
#include "AssertThread.h"

thread_local int g_nFlags;

BOOL WINAPI DllMain(HMODULE hModule, DWORD dwReason, LPVOID lpReserved)
{
	switch (dwReason)
	{
	case DLL_PROCESS_ATTACH:
		break;

	case DLL_PROCESS_DETACH:
		break;
	}

	return TRUE;
}

namespace AssertThread
{

enum THREADINFOCLASS
{
	ThreadQuerySetWin32StartAddress = 9
};

bool IsInThreadpool()
{
	HMODULE hModule = GetModuleHandleW(L"ntdll.dll");

	long (WINAPI *NtQueryInformationThread)(HANDLE, THREADINFOCLASS, PVOID, ULONG, PULONG);
	(PVOID&)NtQueryInformationThread = GetProcAddress(hModule, "NtQueryInformationThread");

	PVOID pStartAddress = nullptr;
	ULONG nLength = 0;
	NtQueryInformationThread(GetCurrentThread(), ThreadQuerySetWin32StartAddress, &pStartAddress, sizeof(PVOID), &nLength);

	if (!pStartAddress || nLength != sizeof(PVOID))
		return false;

	LPCSTR pImageBase = (LPCSTR)hModule;
	DWORD dwImageSize = PIMAGE_NT_HEADERS((LPCSTR)hModule + PIMAGE_DOS_HEADER(hModule)->e_lfanew)->OptionalHeader.SizeOfImage;

	return pStartAddress > pImageBase && pStartAddress < pImageBase + dwImageSize;
}

void AssertThread::SetInThread(ThreadType type, bool bSet)
{
	if ((UINT)type >= ThreadTypeMax)
		return;

	if (bSet)
		g_nFlags |= 1 << type;
	else
		g_nFlags &= ~(1 << type);
}

bool AssertThread::IsInThread(ThreadType type)
{
	if ((UINT)type >= ThreadTypeMax)
		return false;

	if (g_nFlags == 0)
	{
		if (IsInThreadpool())
			g_nFlags |= 1 << Threadpool;
		else
			g_nFlags |= 1 << ThreadTypeMax;
	}

	return (g_nFlags & (1 << type)) != 0;
}

}	// namespace AssertThread
